import os
import logging
from typing import Dict, Optional
from langchain_community.vectorstores import Chroma
from config import Config
from .embedding_manager import ArkEmbeddings

logger = logging.getLogger(__name__)

class ChromaManager:
    """全局ChromaDB实例管理器"""
    
    _instance = None
    _chroma_instances: Dict[str, Chroma] = {}
    _embeddings = None
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(ChromaManager, cls).__new__(cls)
        return cls._instance
    
    def __init__(self):
        if not hasattr(self, '_initialized'):
            self._initialized = True
            self._embeddings = ArkEmbeddings(
                api_key=Config.OPENAI_API_KEY,
                model_name=Config.EMBEDDING_MODEL_NAME
            )
            # 确保向量数据库目录存在
            os.makedirs(Config.VECTOR_DB_PATH, exist_ok=True)
    
    def get_chroma_instance(self, collection_name: str = "default") -> Chroma:
        """获取ChromaDB实例，如果不存在则创建"""
        cache_key = f"{Config.VECTOR_DB_PATH}_{collection_name}"
        
        if cache_key not in self._chroma_instances:
            logger.info(f"创建新的ChromaDB实例: {collection_name}")
            try:
                # 使用ChromaDB客户端直接管理，避免LangChain实例冲突
                import chromadb
                from chromadb.config import Settings
                
                # 创建ChromaDB客户端
                client = chromadb.PersistentClient(
                    path=Config.VECTOR_DB_PATH,
                    settings=Settings(anonymized_telemetry=False)
                )
                
                # 检查集合是否存在
                try:
                    collection = client.get_collection(name=collection_name)
                    logger.info(f"集合 {collection_name} 已存在，使用现有集合")
                except:
                    # 集合不存在，创建新集合
                    collection = client.create_collection(name=collection_name)
                    logger.info(f"创建新集合: {collection_name}")
                
                # 使用LangChain的Chroma包装器
                vector_store = Chroma(
                    client=client,
                    collection_name=collection_name,
                    embedding_function=self._embeddings
                )
                
                self._chroma_instances[cache_key] = vector_store
                logger.info(f"成功加载ChromaDB实例: {collection_name}")
            except Exception as e:
                logger.error(f"加载ChromaDB实例失败: {str(e)}")
                raise
        else:
            logger.debug(f"使用缓存的ChromaDB实例: {collection_name}")
        
        return self._chroma_instances[cache_key]
    
    def create_chroma_instance(self, collection_name: str = "default") -> Chroma:
        """创建新的ChromaDB实例"""
        cache_key = f"{Config.VECTOR_DB_PATH}_{collection_name}"
        
        logger.info(f"创建新的ChromaDB实例: {collection_name}")
        
        # 使用ChromaDB客户端直接管理
        import chromadb
        from chromadb.config import Settings
        
        # 创建ChromaDB客户端
        client = chromadb.PersistentClient(
            path=Config.VECTOR_DB_PATH,
            settings=Settings(anonymized_telemetry=False)
        )
        
        # 删除现有集合（如果存在）
        try:
            client.delete_collection(name=collection_name)
            logger.info(f"删除现有集合: {collection_name}")
        except:
            pass
        
        # 创建新集合
        collection = client.create_collection(name=collection_name)
        logger.info(f"创建新集合: {collection_name}")
        
        # 使用LangChain的Chroma包装器
        vector_store = Chroma(
            client=client,
            collection_name=collection_name,
            embedding_function=self._embeddings
        )
        
        self._chroma_instances[cache_key] = vector_store
        logger.info(f"成功创建ChromaDB实例: {collection_name}")
        
        return vector_store
    
    def clear_cache(self):
        """清除所有缓存的实例"""
        self._chroma_instances.clear()
        logger.info("已清除所有ChromaDB实例缓存")
    
    def get_cached_instances(self) -> Dict[str, str]:
        """获取当前缓存的所有实例信息"""
        return {key: "ChromaDB实例" for key in self._chroma_instances.keys()}

# 全局实例
chroma_manager = ChromaManager() 